import os
import sys
import time
import math
from collections import Counter

import numpy as np
import torch
# 添加yolo模块的路径
yolo_module_path = "/ultralytics/ultralytics"
if os.path.exists(yolo_module_path):
    sys.path.append(yolo_module_path)
from ultralytics import YOLO

# 使能从当前文件定位到 rs2 中的 GStreamerWrapper
_CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
_RS2_DIR = os.path.normpath(os.path.join(_CURRENT_DIR, '..', 'rs2'))
if _RS2_DIR not in sys.path:
    sys.path.append(_RS2_DIR)

from GStreamerWrapper import GStreamerWrapper


class YoloInspector:
    """
    一个YOLO巡检任务处理器，采用空间关联方法完成任务。
    - 识别视野中最主要的目标（以仪表盘为准）作为锚点。
    - 通过计算距离，关联与锚点最近的物体。
    - 在2秒的时间窗口内收集和稳定识别结果。
    - 根据要求格式化并输出最终结果。
    """

    def __init__(self, model_path: str, conf_threshold: float = 0.7):
        """
        初始化巡检器（支持 TensorRT .engine 模型）。
        :param model_path: YOLO TensorRT 引擎文件路径 (.engine)
        :param conf_threshold: 置信度阈值
        """
        # 设备检查
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print(f"正在使用设备: {self.device}")

        # 加载 TensorRT 引擎（Ultralytics 支持直接加载 .engine）
        if not os.path.isfile(model_path) or not model_path.lower().endswith('.engine'):
            raise FileNotFoundError(
                f"模型文件无效，请提供 .engine 文件。当前: '{model_path}'"
            )
        self.model = YOLO(model_path)

        # 状态管理变量
        self.is_inspecting = False
        self.inspection_start_time = None
        self.INSPECTION_WINDOW_S = 2.0
        self.CONF_THRESHOLD = conf_threshold

        # 用于存储2秒窗口内的检测结果
        self.collected_colors = []
        self.collected_numbers = []
        self.collected_states = []

        # 类别名称定义
        self.COLOR_CLASSES = ["red", "orange", "blue", "green", "yellow"]
        self.NUMBER_CLASSES = ["one", "two", "three", "four", "five", "six", "seven", "eight"]
        self.STATE_CLASSES = ["down", "normal", "up"]

        # 标签到中文的映射
        self.color_map = {"red": "红色", "orange": "橙色", "blue": "蓝色", "green": "绿色", "yellow": "黄色"}
        self.number_map = {
            "one": "1号",
            "two": "2号",
            "three": "3号",
            "four": "4号",
            "five": "5号",
            "six": "6号",
            "seven": "7号",
            "eight": "8号",
        }
        self.state_map = {"down": "偏低", "normal": "正常", "up": "偏高"}

    def _get_center(self, box):
        x1, y1, x2, y2 = box
        return (x1 + x2) / 2, (y1 + y2) / 2

    def _get_dominant_object_info(self, boxes, class_names, target_classes):
        dominant_info = None
        max_area = -1
        if not boxes:
            return None
        for i in range(len(boxes)):
            class_id = int(boxes.cls[i].item())
            class_name = class_names.get(class_id)
            if class_name in target_classes:
                box = boxes.xyxy[i].tolist()
                area = (box[2] - box[0]) * (box[3] - box[1])
                if area > max_area:
                    max_area = area
                    dominant_info = {'class_name': class_name, 'box': box}
        return dominant_info

    def _find_closest_object(self, anchor_box, all_boxes, class_names, target_classes):
        anchor_center = self._get_center(anchor_box)
        closest_object_name = None
        min_dist = float('inf')
        if not all_boxes:
            return None
        for i in range(len(all_boxes)):
            class_id = int(all_boxes.cls[i].item())
            class_name = class_names.get(class_id)
            if class_name in target_classes:
                box = all_boxes.xyxy[i].tolist()
                center = self._get_center(box)
                dist = math.dist(anchor_center, center)
                if dist < min_dist:
                    min_dist = dist
                    closest_object_name = class_name
        return closest_object_name

    def process_frame(self, frame: np.ndarray) -> None:
        """处理单帧图像，执行检测并根据状态机管理结果。"""
        results = self.model.predict(frame, verbose=False, conf=self.CONF_THRESHOLD)
        if not results or not results[0].boxes:
            if self.is_inspecting and (time.time() - self.inspection_start_time > self.INSPECTION_WINDOW_S):
                self._finalize_and_print_results()
            return

        result = results[0]
        boxes = result.boxes
        class_names = result.names

        dominant_state_info = self._get_dominant_object_info(boxes, class_names, self.STATE_CLASSES)
        if dominant_state_info is None:
            if self.is_inspecting and (time.time() - self.inspection_start_time > self.INSPECTION_WINDOW_S):
                self._finalize_and_print_results()
            return

        if not self.is_inspecting:
            print("检测到目标，开始2秒钟的稳定识别...")
            self.is_inspecting = True
            self.inspection_start_time = time.time()
            self.collected_colors.clear()
            self.collected_numbers.clear()
            self.collected_states.clear()

        if self.is_inspecting:
            anchor_box = dominant_state_info['box']
            state_name = dominant_state_info['class_name']

            closest_color = self._find_closest_object(anchor_box, boxes, class_names, self.COLOR_CLASSES)
            closest_number = self._find_closest_object(anchor_box, boxes, class_names, self.NUMBER_CLASSES)

            self.collected_states.append(state_name)
            if closest_color:
                self.collected_colors.append(closest_color)
            if closest_number:
                self.collected_numbers.append(closest_number)

            if time.time() - self.inspection_start_time > self.INSPECTION_WINDOW_S:
                self._finalize_and_print_results()

    def _finalize_and_print_results(self) -> None:
        print("2秒识别结束，正在生成最终结果...")
        try:
            if not self.collected_states:
                print("在检测窗口内未能稳定识别出任何仪表盘状态。\n")
                return
            final_state_en = Counter(self.collected_states).most_common(1)[0][0]
            final_color_en = Counter(self.collected_colors).most_common(1)[0][0] if self.collected_colors else None
            final_number_en = Counter(self.collected_numbers).most_common(1)[0][0] if self.collected_numbers else None

            final_color_zh = self.color_map.get(final_color_en, "未知颜色")
            final_number_zh = self.number_map.get(final_number_en, "未知编号")
            final_state_zh = self.state_map.get(final_state_en)

            output_parts = [f"{final_color_zh}区域", final_number_zh, f"仪表盘{final_state_zh}"]
            final_output = " ".join(output_parts)

            print("\n--- 最终巡检结果 ---")
            print(final_output)
            print("-" * 22 + "\n")
        except Exception as e:
            print(f"处理结果时发生错误: {e}\n")
        finally:
            self.is_inspecting = False
            self.inspection_start_time = None
            self.collected_colors.clear()
            self.collected_numbers.clear()
            self.collected_states.clear()


def main():
    # !!!重要!!! 将此路径改为你自己的 TensorRT 引擎文件路径 (.engine)
    # 
    model_path = r"/ultralytics/file/yolo_weight/dashboard-300-epochs-0808.engine"

    try:
        inspector = YoloInspector(model_path, conf_threshold=0.7)
    except Exception as e:
        print(f"错误：无法加载模型，请检查路径 '{model_path}' 是否正确。错误信息: {e}")
        return

    # 初始化 GStreamerWrapper（1920x1080）
    try:
        gs = GStreamerWrapper(width=1920, height=1080)
    except Exception as e:
        print(f"错误：无法初始化 GStreamer 管道。错误信息: {e}")
        return

    print("\nGStreamer 管道已启动，正在等待巡检目标... 按 Ctrl+C 退出。")

    try:
        while True:
            frame = gs.GetFrame()
            if frame is None:
                time.sleep(0.01)
            frame = frame[:,480:1440]
            inspector.process_frame(frame)
            time.sleep(0.01)
    except KeyboardInterrupt:
        print("\n检测已由用户停止。")
    finally:
        try:
            gs.StopThread()
        except Exception:
            pass
        del gs
        print("资源已释放。")


if __name__ == "__main__":
    main()


